#include "lidar_undistortion_2d.h"

Lidar2DUndistortion::Lidar2DUndistortion(){
    __imu_max_mums = 150;
    __imu_cache_buffer = ImuCircularBuffer(__imu_max_mums);    // IMU 100Hz  预留1.5秒数据

    int laser_size = 720;
    __scan_undistortion.angle_min = -M_PI;
    __scan_undistortion.angle_max = M_PI;
    __scan_undistortion.angle_increment = (__scan_undistortion.angle_max - __scan_undistortion.angle_min) / (laser_size - 1);
    __scan_undistortion.range_min = 0.05;
    __scan_undistortion.time_increment = 0.0;
    __scan_undistortion.scan_time = 0.0;    
    __scan_undistortion.ranges.assign(laser_size, std::numeric_limits<float>::infinity());
    __scan_undistortion.intensities.assign(laser_size, std::numeric_limits<float>::infinity());
}

Lidar2DUndistortion::~Lidar2DUndistortion(){

}

// 添加IMU数据
void Lidar2DUndistortion::PushImuData(const sensor_msgs::ImuConstPtr& msg){
    __imu_cache_buffer.push_back(*msg);
}


// 处理雷达数据
void Lidar2DUndistortion::HandleLidarScanData(const sensor_msgs::LaserScanConstPtr& msg, 
    sensor_msgs::LaserScan& output_msg){
    if(__imu_cache_buffer.size() < __imu_max_mums){
        ROS_INFO_STREAM("need more imu data. __imu_cache_buffer size is " 
            << __imu_cache_buffer.size());
        return;
    }


    
    output_msg = __scan_undistortion;
    output_msg.header = msg->header;

    int laser_size = msg->ranges.size();
    ros::Time lidar_timebase = msg->header.stamp;  
    ros::Time lidar_end_timebase = lidar_timebase + ros::Duration(msg->scan_time);
    for(int m=0; m<100; m++){
        if(__imu_cache_buffer.back().header.stamp > lidar_end_timebase){
            break;
        }
        ros::spinOnce();
        ros::Duration(0.01).sleep();
    }

    // 获取IMU对应的起点索引
    int imu_start_idx = -1;
    for(int n=0; n<__imu_cache_buffer.size(); n++){
        if(__imu_cache_buffer[n].header.stamp > lidar_timebase){
            imu_start_idx = n-1;
            break;
        }
    }
    if(imu_start_idx < 0){
        ROS_WARN_STREAM("can not find first timestamp from imu data.");
        return;
    }
    cout << "imu_start_idx: " << imu_start_idx << endl;

    for(int i=0; i<laser_size; i++){
        ros::Time point_timestamp = lidar_timebase + ros::Duration(i * msg->time_increment);
        // 雷达原始坐标系
        float angle = msg->angle_min + i * msg->angle_increment;
        float laser_x = msg->ranges[i] * cos(angle);
        float laser_y = msg->ranges[i] * sin(angle);

        // 计算IMU对应的旋转
        int imu_point_idx = -1;
        for(int n=imu_start_idx; n<__imu_cache_buffer.size(); n++){
            if(__imu_cache_buffer[n].header.stamp > point_timestamp){
                imu_point_idx = n;
                break;
            }
        }
        if(imu_point_idx < 0){
            ROS_WARN_STREAM("can not find point timestamp from imu data. i = "
                << i);
            continue;
        }

        if(imu_start_idx < 0 || imu_point_idx >= __imu_cache_buffer.size()){
            cout << "i: " << i
                << ", imu_start_idx: " << imu_start_idx 
                << ", imu_point_idx: " << imu_point_idx 
                << ", buffer size: " << __imu_cache_buffer.size()
                << endl;
            ROS_WARN_STREAM("out of range.");
            return;
        }

        // 获取球面插值
        auto start_imu = __imu_cache_buffer[imu_start_idx];
        auto end_imu = __imu_cache_buffer[imu_point_idx];
        double t = double((point_timestamp - start_imu.header.stamp).toNSec() / (end_imu.header.stamp - start_imu.header.stamp).toNSec()); 
        if(t < 0){
            ROS_WARN_STREAM("t < 0.0. i= " << i << ", t= " << t);
            continue;
        }

        tf2::Quaternion start_quat;
        tf2::fromMsg(__imu_cache_buffer[imu_start_idx].orientation, start_quat); 
        tf2::Quaternion end_quat;
        tf2::fromMsg(__imu_cache_buffer[imu_point_idx].orientation, end_quat); 
        // 计算当前点的旋转
        tf2::Quaternion quat_interp = start_quat.slerp(end_quat, t);
        // 计算起始点到当前点的旋转
        tf2::Quaternion quat_delta = start_quat.inverse() * quat_interp;

        // 计算旋转后的位姿
        tf2::Vector3 translation(0.0, 0.0, 0.0);
        tf2::Transform transform(quat_delta, translation);
        tf2::Vector3 point_in(laser_x, laser_y, 0.0);
        tf2::Vector3 corrected_point = transform.inverse() * point_in;

        float corrected_range = hypot(corrected_point.x(), corrected_point.y());
        float corrected_angle = atan2(corrected_point.y(), corrected_point.x());
        int index = (int)((corrected_angle - output_msg.angle_min) / output_msg.angle_increment);
        if(index < 0 || index > output_msg.ranges.size()){
            continue;
        }
        output_msg.ranges[index] = min(output_msg.ranges[index], corrected_range);
        output_msg.intensities[index] = msg->intensities[i];
        output_msg.range_max = max(output_msg.angle_max, corrected_range);
    }
}